# table_from_boxes_overlay.py
from typing import List, Tuple, Any
from PIL import Image, ImageDraw, ImageFont
import requests
import time
import sys
import os
# 获取到 config.py 所在的目录（相对于 get_table_cell.py 的路径）
current_dir = os.path.dirname(os.path.abspath(__file__))
config_dir = os.path.abspath(os.path.join(current_dir, '../../..'))  # 返回到 src 目录

# 将 config.py 所在的路径添加到 sys.path
sys.path.append(config_dir)
from config import config
TOL = config.tol
CELL_MODEL_URL = config.cell_model_url
Box = Tuple[float, float, float, float]  # (x1, y1, x2, y2)

# ---------- 与之前相同的聚类 & 构网格函数 ---------- #
def _cluster_lines(coords: List[float], tol: int = TOL) -> List[float]:
    coords_sorted = sorted(coords)
    clusters = [[coords_sorted[0]]]
    for c in coords_sorted[1:]:
        if abs(c - clusters[-1][-1]) <= tol:
            clusters[-1].append(c)
        else:
            clusters.append([c])
    return [sum(g) / len(g) for g in clusters]

def build_table_structure(cell_boxes: List[Box], tol: int = TOL):
    xs, ys = [], []
    for x1, y1, x2, y2 in cell_boxes:
        xs.extend([x1, x2]); ys.extend([y1, y2])

    x_lines, y_lines = _cluster_lines(xs, tol), _cluster_lines(ys, tol)
    n_rows, n_cols = len(y_lines)-1, len(x_lines)-1
    grid: List[List[Any]] = [[None]*n_cols for _ in range(n_rows)]

    for idx, (x1, y1, x2, y2) in enumerate(cell_boxes):
        row_top = max(i for i, y in enumerate(y_lines) if y <= y1+tol)
        row_bot = min(i for i, y in enumerate(y_lines) if y >= y2-tol) - 1
        col_l   = max(j for j, x in enumerate(x_lines) if x <= x1+tol)
        col_r   = min(j for j, x in enumerate(x_lines) if x >= x2-tol) - 1

        rs, cs = row_bot-row_top+1, col_r-col_l+1
        grid[row_top][col_l] = dict(id=idx, rowspan=rs, colspan=cs)
        for r in range(row_top, row_top+rs):
            for c in range(col_l, col_l+cs):
                if (r, c) != (row_top, col_l):
                    grid[r][c] = "SPAN"
    return x_lines, y_lines, grid

def grid_to_html(grid, match_cell_text,header_row: int = 0) -> str:
    html = ["<table>"]
    for r, row in enumerate(grid):
        html.append("<tr>")
        c = 0
        while c < len(row):
            cell = row[c]
            if cell in (None, "SPAN"):
                c += 1; continue
            tag = "th" if r == header_row else "td"
            attrs = []
            if cell["rowspan"] > 1: attrs.append(f'rowspan="{cell["rowspan"]}"')
            if cell["colspan"] > 1: attrs.append(f'colspan="{cell["colspan"]}"')
            attr_str = " "+ " ".join(attrs) if attrs else ""
            # html.append(f'    <{tag}{attr_str}>cell_{cell["id"]}</{tag}>')
            html.append(f'<{tag}{attr_str}>{match_cell_text.get(cell["id"], "")}</{tag}>')
            c += cell["colspan"]
        html.append("</tr>")
    html.append("</table>")
    return "".join(html)

# ---------- 可视化（叠加在原图） ---------- #
def overlay_on_jpg(img_path: str,
                   cell_boxes: List[Box],
                   x_lines: List[float],
                   y_lines: List[float],
                   grid: List[List[Any]]):

    base = Image.open(img_path).convert("RGB")
    font = ImageFont.load_default()

    # ---------- ① 画 OCR box ----------
    img_boxes = base.copy()
    db = ImageDraw.Draw(img_boxes, "RGBA")
    for box in cell_boxes:
        db.rectangle(box, outline=(255, 0, 0, 150), width=2)
    img_boxes.save("./output/debug_boxes.jpg")

    # ---------- ② 画分割线 ----------
    img_lines = base.copy()
    dl = ImageDraw.Draw(img_lines)
    w, h = img_lines.size
    for y in y_lines:
        dl.line([(0, y), (w, y)], fill=(255, 0, 0), width=3)
    for x in x_lines:
        dl.line([(x, 0), (x, h)], fill=(255, 0, 0), width=3)
    img_lines.save("./output/debug_lines.jpg")

    # ---------- ③ 画真实网格 + 合并框 ----------
    img_grid = base.copy()
    dg = ImageDraw.Draw(img_grid)
    for r, row in enumerate(grid):
        for c, cell in enumerate(row):
            x1, y1 = x_lines[c], y_lines[r]
            x2, y2 = x_lines[c + 1], y_lines[r + 1]

            # 单格蓝框
            dg.rectangle([x1, y1, x2, y2], outline=(255, 0, 0), width=1)

            # 合并单元格粗框
            if cell and cell != "SPAN":
                if cell["rowspan"] > 1 or cell["colspan"] > 1:
                    rs, cs = cell["rowspan"], cell["colspan"]
                    dg.rectangle([x1, y1, x_lines[c + cs], y_lines[r + rs]],
                                 outline=(255, 0, 0), width=3)
                dg.text((x1 + 2, y1 + 2), str(cell["id"]),
                        fill=(0, 0, 0), font=font)
    img_grid.save("./output/debug_grid.jpg")

def send_request_with_retry(img, max_retries=1):
    url = f"http://{CELL_MODEL_URL}"
    retries = 0

    while retries <= max_retries:
        try:
            # 发送 POST 请求
            response = requests.post(url, img=img)
            
            # 如果响应状态码是 200（成功）
            if response.status_code == 200:
                response_data = response.json()
                cell_coordinates = response_data['cell_coordinates']
                return cell_coordinates
            
            # 如果响应状态码不是 200，抛出异常
            else:
                raise Exception(f"Unexpected response code: {response.status_code}")

        except (requests.exceptions.RequestException, Exception) as e:
            retries += 1
            print(f"Request failed, retrying {retries}/{max_retries}... Error: {e}")
            
            # 如果达到最大重试次数，抛出异常
            if retries > max_retries:
                raise Exception(f"Request failed after {max_retries} retries.")
            
            # 等待一段时间后重试
            time.sleep(2)  # 你可以根据需要调整重试间隔时间（这里设置为 2 秒）

def table_structure_pre(img,text_cell,img_path=None,tol: int = TOL,visal: str=False):
    # 获取到 re_table 目录的绝对路径
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 将 re_table 所在的目录添加到 sys.path
    sys.path.append(current_dir)
    # from get_table_cell import get_cell_coordinates
    # cell_coordinates = get_cell_coordinates(img)
    cell_coordinates = send_request_with_retry(img)
    dt_boxes = [cell[0] for cell in text_cell]
    from text_cell_match import match_result
    match_re = match_result(dt_boxes,cell_coordinates)
    match_cell_text = {k:''.join([text_cell[i][1] for i in v]) for k,v in match_re.items()}
    x_lines, y_lines, grid = build_table_structure(cell_coordinates, tol=5)
    if visal and image_path:
        overlay_on_jpg(img_path,cell_coordinates,x_lines,y_lines,grid)
    html = grid_to_html(grid,match_cell_text, header_row=0)
    return html
# ---------------------- DEMO ---------------------- #
if __name__ == "__main__":
    # ←←把下方两处替换成你的数据和图片路径→→
    image_path = "/home/fengjie/GITEE/pdf2md_without_ocr/inputs/1185e552f89d162364ff434d30006bbf1d9f922d0dad48743efeb7064ef78ae0.jpg"      # 原始 JPG 表格
    

    html = table_structure_pre(image_path,image_path)
    print(html)
    # with open("./output/table.md", "w", encoding="utf-8") as f:
    #     f.write(html)
    print("⎋ 调试图已输出：debug_boxes.jpg / debug_lines.jpg / debug_grid.jpg")
    print("⎋ HTML 保存为 table.md")